{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numbers\n",
    "from PIL import Image, ImageMath\n",
    "import os\n",
    "import os.path\n",
    "import numpy as np\n",
    "import struct\n",
    "import math\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.use('TkAgg')\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import matplotlib.cm as cm\n",
    "\n",
    "import time\n",
    "import itertools\n",
    "\n",
    "import pcl\n",
    "import timeit\n",
    "import png\n",
    "\n",
    "import multiprocessing\n",
    "\n",
    "from plyfile import PlyData, PlyElement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def axisEqual3D(ax):\n",
    "    extents = np.array([getattr(ax, 'get_{}lim'.format(dim))() for dim in 'xyz'])\n",
    "    sz = extents[:, 1] - extents[:, 0]\n",
    "    centers = np.mean(extents, axis=1)\n",
    "    maxsize = max(abs(sz))\n",
    "    r = maxsize / 2\n",
    "    for ctr, dim in zip(centers, 'xyz'):\n",
    "        getattr(ax, 'set_{}lim'.format(dim))(ctr - r, ctr + r)\n",
    "\n",
    "\n",
    "def plot_pc(pc_np, z_cutoff=1000, birds_view=False, color='height', size=0.3, ax=None, cmap=cm.jet, is_equal_axes=True):\n",
    "    # remove large z points\n",
    "    valid_index = pc_np[:, 0] < z_cutoff\n",
    "    pc_np = pc_np[valid_index, :]\n",
    "\n",
    "    if ax is None:\n",
    "        fig = plt.figure(figsize=(9, 9))\n",
    "        ax = Axes3D(fig)\n",
    "    if type(color)==str and color == 'height':\n",
    "        c = pc_np[:, 2]\n",
    "        ax.scatter(pc_np[:, 0], pc_np[:, 1], pc_np[:, 2], s=size, c=c, cmap=cmap, edgecolors='none')\n",
    "    elif type(color)==str and color == 'reflectance':\n",
    "        assert False\n",
    "    elif type(color) == np.ndarray:\n",
    "        ax.scatter(pc_np[:, 0], pc_np[:, 1], pc_np[:, 2], s=size, c=color, cmap=cmap, edgecolors='none')\n",
    "    else:\n",
    "        ax.scatter(pc_np[:, 0], pc_np[:, 1], pc_np[:, 2], s=size, c=color, edgecolors='none')\n",
    "\n",
    "    if is_equal_axes:\n",
    "        axisEqual3D(ax)\n",
    "    if True == birds_view:\n",
    "        ax.view_init(elev=0, azim=-90)\n",
    "    else:\n",
    "        ax.view_init(elev=-45, azim=-90)\n",
    "    # ax.invert_yaxis()\n",
    "\n",
    "    return ax\n",
    "\n",
    "# get surface normal\n",
    "def Surface_normals(cloud):\n",
    "    ne = cloud.make_NormalEstimation()\n",
    "    tree = cloud.make_kdtree()\n",
    "    ne.set_SearchMethod(tree)\n",
    "#     ne.set_RadiusSearch(2)\n",
    "    ne.set_KSearch(9)\n",
    "    cloud_normals = ne.compute()\n",
    "    return cloud_normals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_root = '/ssd/jiaxin/TSF_datasets/3DMatch_eval'\n",
    "output_root = '/ssd/jiaxin/TSF_datasets/3DMatch_eval_npy'\n",
    "scene_list = ['7-scenes-redkitchen', 'sun3d-home_at-home_at_scan1_2013_jan_1', \n",
    "              'sun3d-home_md-home_md_scan9_2012_sep_30', 'sun3d-hotel_uc-scan3',\n",
    "              'sun3d-hotel_umd-maryland_hotel1', 'sun3d-hotel_umd-maryland_hotel3', \n",
    "              'sun3d-mit_76_studyroom-76-1studyroom2', 'sun3d-mit_lab_hj-lab_hj_tea_nov_2_2012_scan1_erika']\n",
    "downsample_size = 16384\n",
    "\n",
    "for scene in scene_list:\n",
    "    ply_filename_list = os.listdir(os.path.join(dataset_root, scene))\n",
    "    \n",
    "    npy_folder = os.path.join(output_root, scene)\n",
    "    if not os.path.isdir(npy_folder):\n",
    "        os.makedirs(npy_folder)\n",
    "    \n",
    "    for ply_filename in ply_filename_list:\n",
    "        plydata = PlyData.read(os.path.join(dataset_root, scene, ply_filename))\n",
    "        ply_x = plydata['vertex']['x']\n",
    "        ply_y = plydata['vertex']['y']\n",
    "        ply_z = plydata['vertex']['z']\n",
    "        pc_np = np.stack((ply_x, ply_y, ply_z), axis=1)\n",
    "        \n",
    "        # downsample\n",
    "        if pc_np.shape[0] > downsample_size:\n",
    "            choice_idx = np.random.choice(pc_np.shape[0], downsample_size, replace=False)\n",
    "            pc_np = pc_np[choice_idx, :]\n",
    "            \n",
    "        # compute surface normal\n",
    "        cloud = pcl.PointCloud(pc_np.astype(np.float32))    \n",
    "        sn = Surface_normals(cloud)\n",
    "        sn_np = np.asarray(sn.to_array(), dtype=np.float32)  # Nx4, nx,ny,nz,curvature\n",
    "\n",
    "        output_np = np.concatenate((pc_np.astype(np.float32), sn_np), axis=1)  # Nx7\n",
    "        \n",
    "        npy_filename = ply_filename[0:-4] + '.npy'\n",
    "        npy_file = os.path.join(npy_folder, npy_filename)\n",
    "        np.save(npy_file, output_np)\n",
    "#         break;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
